from naie.datasets import get_data_reference
from naie.context import Context
import moxing as mox
import zipfile

# load model state dict
data_reference = get_data_reference(dataset="pretrain",
                                    dataset_entity='resnet50_best')

for file_path in data_reference.get_files_paths():
    mox.file.copy(file_path, './cache/' + file_path.split('/')[-1])

# load data
data_reference = get_data_reference(dataset="DatasetService",
                                    dataset_entity='image_dataset')

for file_path in data_reference.get_files_paths():
    mox.file.copy(file_path, './cache/' + file_path.split('/')[-1])

zip_file = zipfile.ZipFile('./cache/data.zip')
zip_list = zip_file.namelist()
for f in zip_list:
    zip_file.extract(f, './cache/')
